#pragma once
#include <cuda_runtime.h>
#include <Utilities/glm_ext.h>

namespace PhysLeo {

/**
 * FEM StVK material.
 * \f$ E=\frac{F^TF-I}{2} \f$,
 * \f$ W(E) = \mu(E:E)+\frac{\lambda}{2}tr^2(E) \f$
 */
template<typename T>
class FemStvk {
public:
    /**
    * return the energy density according to deformation gradient
    * @param[in] F  a glm 3x3 matrix type represents the deformation gradient.
    * @param[in] ptr_material  a pointer to vertex specific material data
    * @return a T value, the energy density.
    */
    __host__ __device__ static T energyDensity(glm::tmat3x3<T> F, T* ptr_material)
    {
        auto lambda = ptr_material[0];
        auto mu = ptr_material[1];

        glm::tmat3x3<T> identity(1.0f);

        //e = 0.5(F^T*F-I)
        glm::tmat3x3<T> strain = 0.5f*(glm::transpose(F)*F - identity);

        // sum = e:e
        T sum = 0;
        for (auto i = 0; i < 3; ++i)
            for (auto j = 0; j < 3; ++j)
                sum += strain[i][j] * strain[i][j];

        //trait = tr(e)
        T trait = strain[0][0] + strain[1][1] + strain[2][2];

        // W = mu*e:e+lambda/2*tr(e)^2
        return mu * sum + 0.5f*lambda*trait*trait;
    }

    /**
    * return the first piola stress according to deformation gradient.
    * @param[in] F  a glm 3x3 matrix type represents the deformation gradient.
    * @param[in] ptr_material  a pointer to vertex specific material data
    * @return a glm 3x3 matrix type represents the first piola stress.
    */
    __host__ __device__ static glm::tmat3x3<T> firstPiolaStress(glm::tmat3x3<T> F, T* ptr_material)
    {
        auto lambda = ptr_material[0];
        auto mu = ptr_material[1];

        glm::tmat3x3<T> identity(1.0f);

        //e = 0.5(F^T*F-I)
        glm::tmat3x3<T> e = static_cast<T>(0.5)*(glm::transpose(F)*F - identity);

        //P = F(2*mu*e+lambda*tr(e)*I)
        return F * (2.0f*mu*e + lambda * glm::trace(e)*identity);
    }
};

}